package examples.DamenF where

import frege.lib.ForkJoin (par)

native      currentTimeMillis java.lang.System.currentTimeMillis :: () -> IO Long
pure native formatLocale java.lang.String.format :: Maybe JLocale -> String -> Float -> String
data JLocale = pure native java.util.Locale

format s f = formatLocale Nothing s f

main [algo, arg]
    | Right n <- arg.int = do
        start <- currentTimeMillis ()
        print n
        print " Damen "
        case algo of
            "f" -> print (count 0 n ok0 0 0)
            "l" -> print (solutions n)
            "lp" -> print (psolutions n)
            "a" -> print (arraysolutions n)
            "h" -> print (length $ allSolutions n n)
            _   -> print "no "
        print " Solutions"
        stop <- currentTimeMillis ()
        println (" in " ++ format "%.3f" ((stop-start).float / 1e3f) ++ "s.")

main _ = println "usage: java DamenF {f|l|a|lp|h} n"

-- list based solution
fromto i j xs   -- tail recursive version of i..j
    | i <= j = fromto i (j-1) (j:xs)
    | otherwise = xs


sol !ns i row | i>0 = fold (+) 0 [ b |
                        a <- ns,
                        check a row,
                        !b = sol ns (i-1) (a:row)
                    ]
             | otherwise = 1

check n xs  = let
        checkd !y (a:bs) !diff =
            if y != a && y != a+diff && y != a-diff
                then checkd y bs (diff+1)
                else false
        checkd y [] diff     = true
    in checkd n xs 1

solutions  n =  sol (fromto 1 n []) n []

-- n parallel threads
-- the i-th thread computes the number of solutions
-- for the case that a queen stands in column i on row 1
psolutions n = fold (+) 0 (pmap psol ns)
    where
        ns = fromto 1 n []
        psol i = sol (filter (i!=) ns) (n-1) [i]
        pmap f [] = []
        pmap f (x:xs) = a `par` as `par` (a:as)
            where
                a  = f x
                as = pmap f xs

-- functional solution
ok0 :: Int -> Int -> Bool
ok0 _ _ = true
ok1 :: (Int -> Int -> Bool) -> Int -> Int -> Int -> Int -> Bool
ok1 f r1 c1 r2 c2 = c1 != c2 && (c1-(r2-r1)) != c2 && (c1+(r2-r1)) != c2 && f r2 c2


count res n ok r c =
    if       r == n then res+1
    else if  c == n then res
    else if  ok r c then count (count res n (ok1 ok r c) (r+1) 0) n ok r (c+1)
    else                 count res n ok r (c+1)


-- solution with integer array
arraysolutions n = ST.run (arraysol n)


arraysol n = newArray n >>= arrayloop 0 0 where
    arrayloop i !sols !arr
        | i == n = arrayloop (i-1) (sols+1) arr
        | i >= 0 = do  q <- getElemAt arr i
                       q <- readonly  (unattacked (q+1) 0) arr
                       if q > n
                           then let next = i-1 in setElemAt arr i 0 >> arrayloop next sols arr
                           else let next = i+1 in setElemAt arr i q >> arrayloop next sols arr
        | otherwise = return sols
        where unattacked :: Int -> Int -> JArray Int ->Int
              unattacked !q !j !arr
                | q > n = q
                | j == i = q
                | otherwise = case elemAt arr j of
                    0xdeadbeef = 0  -- force aj to be native
                    !aj | aj == q || aj-q == i-j || q-aj == i-j = unattacked (q+1) 0 arr
                        | otherwise = unattacked q (j+1) arr



-- solution by Oystein Kolsrud
-- https://www.youtube.com/watch?v=I2tMmsZC1ZU
okToAdd :: Int -> [Int] -> Bool
okToAdd q qs = all (okToAddDirection q qs) [succ, pred, id]
    where
        okToAddDirection q qs f = and $ zipWith (/=) (tail (iterate f q)) qs

extendSolution n qs = map (\q -> q:qs) $ filter (\q -> okToAdd q qs) [1..n]

allSolutions !n 0 = [[]]
allSolutions !n k = concatMap (extendSolution n) (allSolutions n (k-1))




